!!!! OEGM Toolbox - (12/21/2022) !!!!
!! Subroutines used to solve the Overlapping Endogenous Grid Method
!! OEGM: calculates an upper envelop using overlapping endogneous grids

!!!! Table of Contents !!!!
!! 1.) Endogenous Grids, subroutine to form endogenous grids and overlapping regions
!! 2.) OEGM Interpolation, subroutine to solve for decision rules by interpolating onto endogenous grids

module OEGM_TB

use BCI_Par
use Fortran_TB

implicit none

contains

!!Introduction to Code
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
!! 1.) Endogenous Grids
    
subroutine endogenous_grids(beta, rate, income, agrid, DV, n_ogrid, ilow_ogrid, ihigh_ogrid, EnGrid)

!!!! Variables and Arrays !!!!
real(rp), intent(in) :: beta, rate, income
real(rp), dimension(na), intent(in) :: agrid, DV
integer, intent(out) :: n_ogrid
integer, dimension(no), intent(out) :: ilow_ogrid, ihigh_ogrid
real(rp), dimension(na), intent(out) :: EnGrid
integer :: a2
real(rp) :: c_hat

EnGrid(:) = -agrid(na)
n_ogrid = 1
ilow_ogrid(1) = nd
do a2 = nd,na

    !!!! Form Endogenous Grid !!!!
    c_hat = ( beta*(one + rate)*DV(a2) )**(-one/sigma)
    EnGrid(a2) = c_hat + agrid(a2)/(one+rate) - income - y_bar
    
    !!!! Form Overlapping Regions !!!!
    if ( EnGrid(a2) <= EnGrid(a2-1) ) then
        ihigh_ogrid(n_ogrid) = a2 - 1
        n_ogrid = n_ogrid + 1
        ilow_ogrid(n_ogrid) = a2
    end if
    
    !!!! Check for Errors !!!!
    if ( beta*(one + rate)*DV(a2) <= zero ) then
        print *, "WARNING, forming Endogenous Grids"
    end if
end do
ihigh_ogrid(n_ogrid) = na

if ( n_ogrid > no ) then
    print *, "WARNING, too many overlapping endogenous grids"
end if

end subroutine endogenous_grids

!! 1.) Endogenous Grids
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
!! 2.) OEGM Interpolation

subroutine oegm_interpolation(n_oegm, beta, chi, a_val, income, rate, ilow, ihigh, EnGrid, agrid, EV, maxV, maxG, maxC, maxQ)

!!!! Variables and Arrays !!!!
integer, intent(in) :: n_oegm
integer, dimension(no), intent(in) :: ilow, ihigh
real(rp), intent(in) :: beta, chi, a_val, income, rate
real(rp), dimension(na), intent(in) :: EnGrid, agrid, EV
real(rp), intent(inout) :: maxV, maxG, maxC, maxQ
integer :: x, ia, io_low, io_high
real(rp) :: omega, Value, consumption, saving

!!!! Decision Rules in Overlapping Regions !!!!
do x = 1,n_oegm
    io_low = ilow(x)
    io_high = ihigh(x)
    if ( io_low == io_high ) then
        ia = io_low
        omega = one
    else
        call interpolation(EnGrid(:), a_val, io_low, io_high, na, ia, omega)
    end if
    
    saving = omega*agrid(ia) + (one-omega)*agrid(ia+1)
    consumption = a_val + income + y_bar - saving/(one+rate) 
    if ( consumption >= zero ) then
        Value = (one/(one-sigma))*(( consumption )**(one-sigma)) - chi + beta*( omega*EV(ia) + (one-omega)*EV(ia+1) )
    else
        Value = disutility
    end if
    
    if ( Value > maxV ) then
        maxV = Value
        maxC = consumption
        maxG = saving
        maxQ = one/(one+rate)
    end if
end do

end subroutine oegm_interpolation

!! 2.) OEGM Decision Rules 
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

end module OEGM_TB